
from typing import List
from sentence_transformers import SentenceTransformer
import faiss

from multihop.core.docment import Document, RetrievalResult, logger


class VectorDatabase:
    """FAISS向量数据库管理器"""

    def __init__(self, embedding_model_name: str = "all-MiniLM-L6-v2"):
        self.embedding_model = SentenceTransformer(embedding_model_name)
        self.dimension = self.embedding_model.get_sentence_embedding_dimension()
        self.index = faiss.IndexFlatIP(self.dimension)  # 内积索引
        self.documents: List[Document] = []

    def add_documents(self, documents: List[Document]):
        """添加文档到向量数据库"""
        if not documents:
            return

        # 提取文档内容并生成嵌入
        texts = [doc.content for doc in documents]
        embeddings = self.embedding_model.encode(texts, normalize_embeddings=True)

        # 添加到FAISS索引
        self.index.add(embeddings.astype('float32'))
        self.documents.extend(documents)

        logger.info(f"已添加 {len(documents)} 个文档到向量数据库")

    def search(self, query: str, k: int = 5) -> List[RetrievalResult]:
        """搜索相似文档"""
        if self.index.ntotal == 0:
            return []

        # 生成查询嵌入
        query_embedding = self.embedding_model.encode([query], normalize_embeddings=True)

        # 搜索最相似的文档
        scores, indices = self.index.search(query_embedding.astype('float32'), min(k, self.index.ntotal))

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx != -1:  # FAISS返回-1表示无效索引
                results.append(RetrievalResult(
                    document=self.documents[idx],
                    score=float(score),
                    hop_level=0  # 初始为0，后续会更新
                ))

        return results